#!/usr/bin/env python
import time
import py_trees
import py_trees.console as console


def pre_tick_handler(behaviour_tree: py_trees.trees.BehaviourTree) -> None:
    """Generate a simple pre-tick banner printing to stdout."""
    print("\n--------- Run %s ---------\n" % behaviour_tree.count)


class SuccessEveryN(py_trees.behaviours.SuccessEveryN):
    """Add a blackboard counter to :class:`~py_trees.behaviours.SuccessEveryN`."""

    def __init__(self) -> None:
        """Set up the blackboard."""
        super().__init__(name="EveryN", n=5)
        self.blackboard = self.attach_blackboard_client(name=self.name)
        self.blackboard.register_key("count", access=py_trees.common.Access.WRITE)

    def update(self) -> py_trees.common.Status:
        """
        Run the :class:`~py_trees.behaviours.SuccessEveryN` update and write to blackboard.
        Returns:
            Returns the :class:`~py_trees.common.Status` of :class:`~py_trees.behaviours.SuccessEveryN`
        """
        status = super().update()
        self.blackboard.count = self.count
        return status


class PeriodicSuccess(py_trees.behaviours.Periodic):
    """Write the period from :class:`~py_trees.behaviours.Periodic` as a variable on the blackboard."""

    def __init__(self) -> None:
        """Initialise :class:`~py_trees.behaviours.Periodic` and set up the blackboard."""
        super().__init__(name="Periodic", n=3)
        self.blackboard = self.attach_blackboard_client(name=self.name)
        self.blackboard.register_key("period", access=py_trees.common.Access.WRITE)

    def update(self) -> py_trees.common.Status:
        """
        Run the :class:`~py_trees.behaviours.Periodic` update and write to blackboard.
        Returns:
            Returns the :class:`~py_trees.common.Status` of :class:`~py_trees.behaviours.Periodic`
        """
        status = super().update()
        self.blackboard.period = self.period
        return status


class Finisher(py_trees.behaviour.Behaviour):
    """
    Gathers blackboard data from other behaviours and prints a summary.
    To be used at the end of the tree run.
    """

    def __init__(self) -> None:
        """Set up the blackboard."""
        super().__init__(name="Finisher")
        self.blackboard = self.attach_blackboard_client(name=self.name)
        self.blackboard.register_key("count", access=py_trees.common.Access.READ)
        self.blackboard.register_key("period", access=py_trees.common.Access.READ)

    def update(self) -> py_trees.common.Status:
        """
        Fetch blackboard variables and print a summary.
        Returns:
            Always returns :data:`~py_trees.common.Status.SUCCESS`.
        """
        print(console.green + "---------------------------" + console.reset)
        print(console.bold + "        Finisher" + console.reset)
        print(console.green + "  Count : {}".format(self.blackboard.count) + console.reset)
        print(console.green + "  Period: {}".format(self.blackboard.period) + console.reset)
        print(console.green + "---------------------------" + console.reset)
        return py_trees.common.Status.SUCCESS


def create_tree() -> py_trees.behaviour.Behaviour:
    """
    Create the root behaviour and it's subtree.
    Returns:
        the root behaviour
    """
    every_n_success = SuccessEveryN()
    sequence = py_trees.composites.Sequence(name="Sequence", memory=True)
    guard = py_trees.behaviours.Success("Guard")
    periodic_success = PeriodicSuccess()
    finisher = Finisher()
    sequence.add_child(guard)
    sequence.add_child(periodic_success)
    sequence.add_child(finisher)
    idle = py_trees.behaviours.Success("Idle")
    root = py_trees.composites.Selector(name="Demo Tree", memory=False)
    root.add_child(every_n_success)
    root.add_child(sequence)
    root.add_child(idle)
    return root


##############################################################################
# Main
##############################################################################


def main() -> None:
    """Entry point for the demo script."""

    py_trees.logging.level = py_trees.logging.Level.DEBUG
    tree = create_tree()
    py_trees.display.render_dot_tree(tree, with_blackboard_variables=True)

    ####################
    # Tree Stewardship
    ####################
    py_trees.blackboard.Blackboard.enable_activity_stream(100)
    behaviour_tree = py_trees.trees.BehaviourTree(tree)
    behaviour_tree.add_pre_tick_handler(pre_tick_handler)
    # visitors访客调试
    behaviour_tree.visitors.append(py_trees.visitors.DebugVisitor())
    """每一次tick都会调用这个函数
    打印行为树
        状态
        黑板变量
        黑板变量读写记录
    """
    behaviour_tree.visitors.append(
        py_trees.visitors.DisplaySnapshotVisitor(display_blackboard=True, display_activity_stream=True)
    )
    behaviour_tree.setup(timeout=15)

    ####################
    # Tick Tock
    ####################
    interactive = False
    if interactive:
        py_trees.console.read_single_keypress()
    while True:
        try:
            behaviour_tree.tick()
            if interactive:
                py_trees.console.read_single_keypress()
            else:
                time.sleep(0.5)
        except KeyboardInterrupt:
            break
    print("\n")


if __name__ == "__main__":
    main()
